"""
多信号 SVM 策略
使用 RSRS、OC_DIFF、RET 三个因子，通过 SVM 分类模型预测收益率区间
"""
import numpy as np
import pandas as pd
from sklearn.svm import SVC
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from .base_strategy import BaseStrategy


class MultiSignalSVMStrategy(BaseStrategy):
    strategy_name = "multi_signal_svm"
    
    def __init__(self, api, symbol_info, market_period):
        super().__init__(api, symbol_info, market_period)
        for symbol in symbol_info.keys():
            # 存储三个因子的历史序列
            self.all_info[symbol]["rsrs_slopes"] = []
            self.all_info[symbol]["corr_series"] = []
            self.all_info[symbol]["ret_series"] = []
            
            # 存储标准化后的因子得分序列
            self.all_info[symbol]["rsrs_score_series"] = []
            self.all_info[symbol]["oc_score_series"] = []
            self.all_info[symbol]["ret_score_series"] = []
            
            # 存储收益率序列
            self.all_info[symbol]["return_series"] = []
            
            # 模型相关参数
            self.all_info[symbol]["lookback_period"] = symbol_info[symbol].get("lookback_period", 10)   # 特征窗口（降低到10）
            self.all_info[symbol]["train_period"] = symbol_info[symbol].get("train_period", 200)       # 训练窗口
            self.all_info[symbol]["retrain_interval"] = symbol_info[symbol].get("retrain_interval", 30)  # 重训练间隔（天）
            self.all_info[symbol]["max_series_length"] = symbol_info[symbol].get("max_series_length", 800)
            
            # 收益率离散化参数
            self.all_info[symbol]["return_bins"] = symbol_info[symbol].get("return_bins", 10)   # 分成10个区间（-5到5）
            self.all_info[symbol]["return_scale"] = symbol_info[symbol].get("return_scale", 50.0)  # 收益率缩放系数（降低到50）
            
            # SVM 模型和标准化器
            self.all_info[symbol]["model"] = None
            self.all_info[symbol]["scaler"] = StandardScaler()
            self.all_info[symbol]["last_train_index"] = -1  # 记录上次训练的索引
    
    def calculate_signal(self, symbol, current_pos, kline):
        """计算交易信号 - 基于 SVM 预测的收益率区间"""
        # 更新三个因子的原始值和标准化得分
        self._update_factor_series(symbol, kline)
        
        # 更新收益率序列
        self._update_return_series(symbol, kline)
        
        # 获取参数
        lookback_period = self.all_info[symbol]["lookback_period"]
        train_period = self.all_info[symbol]["train_period"]
        
        # 检查是否有足够的数据进行训练
        min_data_length = lookback_period + train_period
        rsrs_scores = self.all_info[symbol]["rsrs_score_series"]
        
        if len(rsrs_scores) < min_data_length:
            return 0, 0.0
        
        # 训练模型并预测
        predicted_class = self._train_and_predict(symbol)
        
        if predicted_class is None:
            return 0, 0.0
        
        # 将预测的类别转换为仓位信号
        discretize_method = self.all_info[symbol].get("discretize_method", "quantile")
        
        if discretize_method in ["quantile", "adaptive", "threshold"]:
            # 新方法：类别范围 -2 到 2，除以 2 得到 -1 到 1 的仓位
            target_position = predicted_class / 2.0
        else:
            # 原方法：类别范围 -5 到 5，除以 5 得到 -1 到 1 的仓位
            target_position = predicted_class / 5.0
        
        # 限制在 [-1, 1] 范围内
        target_position = np.clip(target_position, -1.0, 1.0)
        
        # DEBUG: 打印前几次的预测
        if len(rsrs_scores) <= min_data_length + 50:
            print(f"Bar {len(rsrs_scores)}: predicted_class={predicted_class}, target_pos={target_position:.3f}")
        
        return target_position, predicted_class
    
    def _update_factor_series(self, symbol, kline):
        """更新三个因子的原始值和标准化得分"""
        # 计算 RSRS 因子
        rsrs_score = self._calculate_rsrs(symbol, kline)
        if rsrs_score is not None:
            self.all_info[symbol]["rsrs_score_series"].append(rsrs_score)
            if len(self.all_info[symbol]["rsrs_score_series"]) > self.all_info[symbol]["max_series_length"]:
                self.all_info[symbol]["rsrs_score_series"].pop(0)
        
        # 计算 OC 相关性因子
        oc_score = self._calculate_oc_corr(symbol, kline)
        if oc_score is not None:
            self.all_info[symbol]["oc_score_series"].append(oc_score)
            if len(self.all_info[symbol]["oc_score_series"]) > self.all_info[symbol]["max_series_length"]:
                self.all_info[symbol]["oc_score_series"].pop(0)
        
        # 计算 RET 因子
        ret_score = self._calculate_ret(symbol, kline)
        if ret_score is not None:
            self.all_info[symbol]["ret_score_series"].append(ret_score)
            if len(self.all_info[symbol]["ret_score_series"]) > self.all_info[symbol]["max_series_length"]:
                self.all_info[symbol]["ret_score_series"].pop(0)
    
    def _update_return_series(self, symbol, kline):
        """更新收益率序列"""
        close_prices = kline.close.values
        
        if len(close_prices) < 2:
            return
        
        # 计算收益率
        returns = np.diff(close_prices) / close_prices[:-1]
        
        # 更新收益率序列
        current_len = len(returns)
        stored_len = len(self.all_info[symbol]["return_series"])
        
        if stored_len < current_len:
            # 批量添加缺失的收益率
            self.all_info[symbol]["return_series"].extend(returns[stored_len:].tolist())
            
            # 限制序列长度
            max_length = self.all_info[symbol]["max_series_length"]
            if len(self.all_info[symbol]["return_series"]) > max_length:
                self.all_info[symbol]["return_series"] = self.all_info[symbol]["return_series"][-max_length:]

    
    def _discretize_return(self, return_value, return_scale, return_bins):
        """
        将收益率离散化为整数类别
        
        方法选择（通过 symbol_info 配置）：
        - 'linear': 线性缩放（默认）
        - 'quantile': 分位数离散化（推荐）
        - 'threshold': 阈值离散化
        - 'adaptive': 自适应阈值
        """
        discretize_method = self.all_info[list(self.all_info.keys())[0]].get("discretize_method", "quantile")
        
        if discretize_method == "linear":
            # 原方法：线性缩放
            scaled_return = return_value * return_scale
            discretized = int(np.clip(np.round(scaled_return), -5, 5))
            return discretized
        
        elif discretize_method == "quantile":
            # 分位数方法：基于历史收益率分布
            # 需要在训练时计算分位数阈值
            return None  # 在 _train_and_predict 中处理
        
        elif discretize_method == "threshold":
            # 阈值方法：使用固定百分比
            threshold_pct = self.all_info[list(self.all_info.keys())[0]].get("threshold_pct", 0.005)  # 0.5%
            
            if return_value > 2 * threshold_pct:
                return 2  # 强做多
            elif return_value > threshold_pct:
                return 1  # 做多
            elif return_value < -2 * threshold_pct:
                return -2  # 强做空
            elif return_value < -threshold_pct:
                return -1  # 做空
            else:
                return 0  # 观望
        
        elif discretize_method == "adaptive":
            # 自适应方法：基于标准差
            return None  # 在 _train_and_predict 中处理
        
        else:
            # 默认使用线性方法
            scaled_return = return_value * return_scale
            discretized = int(np.clip(np.round(scaled_return), -5, 5))
            return discretized
    
    def _train_and_predict(self, symbol):
        """训练 SVM 模型并预测下一期收益率区间（定期重训练）"""
        lookback_period = self.all_info[symbol]["lookback_period"]
        train_period = self.all_info[symbol]["train_period"]
        retrain_interval = self.all_info[symbol]["retrain_interval"]
        return_scale = self.all_info[symbol]["return_scale"]
        return_bins = self.all_info[symbol]["return_bins"]
        discretize_method = self.all_info[symbol].get("discretize_method", "quantile")
        
        # 获取三个因子的得分序列
        rsrs_scores = self.all_info[symbol]["rsrs_score_series"]
        oc_scores = self.all_info[symbol]["oc_score_series"]
        ret_scores = self.all_info[symbol]["ret_score_series"]
        return_series = self.all_info[symbol]["return_series"]
        
        # 检查数据长度 - 确保所有序列长度一致
        min_series_len = min(len(rsrs_scores), len(oc_scores), len(ret_scores))
        
        # 收益率序列可能比因子序列短1（因为收益率需要两个价格点）
        # 确保有足够的数据进行训练
        min_data_required = lookback_period + train_period
        if min_series_len < min_data_required or len(return_series) < min_data_required:
            return None
        
        # 使用较短序列的长度作为基准
        available_len = min(min_series_len, len(return_series))
        
        # 判断是否需要重新训练模型
        current_index = available_len
        last_train_index = self.all_info[symbol]["last_train_index"]
        need_retrain = (self.all_info[symbol]["model"] is None or 
                       current_index - last_train_index >= retrain_interval)
        
        if need_retrain:
            # 构建训练数据
            X_train_list = []
            y_train_list = []
            train_returns = []  # 用于计算分位数
            
            # 滚动窗口构建训练样本
            # 确保索引不会越界
            start_idx = available_len - train_period - lookback_period
            end_idx = available_len - lookback_period
            
            for i in range(start_idx, end_idx):
                # 检查索引是否有效
                if i < 0 or i + lookback_period >= len(return_series):
                    continue
                
                # 获取前 lookback_period 期的三个因子得分作为特征
                features = []
                features.extend(rsrs_scores[i:i+lookback_period])
                features.extend(oc_scores[i:i+lookback_period])
                features.extend(ret_scores[i:i+lookback_period])
                
                X_train_list.append(features)
                
                # 目标是下一期收益率
                next_return = return_series[i + lookback_period]
                train_returns.append(next_return)
            
            # 检查是否有足够的训练样本
            if len(X_train_list) < 10:  # 至少需要10个样本
                print(f"警告：训练样本不足 ({len(X_train_list)} 个)，跳过训练")
                return None
            
            # 根据离散化方法处理标签
            if discretize_method == "quantile":
                # 分位数离散化：确保每个类别样本数相对均衡
                n_classes = 5  # 5个类别：-2, -1, 0, 1, 2
                train_returns_array = np.array(train_returns)
                quantiles = np.linspace(0, 1, n_classes + 1)
                bins = np.quantile(train_returns_array, quantiles)
                bins[0] = -np.inf  # 确保包含所有值
                bins[-1] = np.inf
                
                # 离散化
                y_train_list = np.digitize(train_returns_array, bins) - (n_classes // 2 + 1)
                y_train_list = np.clip(y_train_list, -2, 2)
                
                # 保存分位数阈值供预测使用
                self.all_info[symbol]["quantile_bins"] = bins
                
                print(f"分位数阈值: {bins}")
                print(f"标签分布: {np.bincount(y_train_list + 2)}")  # +2 使索引从0开始
                
            elif discretize_method == "adaptive":
                # 自适应阈值：基于训练集的标准差
                train_returns_array = np.array(train_returns)
                mean_ret = train_returns_array.mean()
                std_ret = train_returns_array.std()
                
                # 使用 0.5 和 1.0 倍标准差作为阈值
                threshold_1 = 0.5 * std_ret
                threshold_2 = 1.0 * std_ret
                
                y_train_list = np.zeros(len(train_returns))
                y_train_list[train_returns_array > mean_ret + threshold_2] = 2
                y_train_list[train_returns_array > mean_ret + threshold_1] = 1
                y_train_list[train_returns_array < mean_ret - threshold_2] = -2
                y_train_list[train_returns_array < mean_ret - threshold_1] = -1
                
                # 保存阈值供预测使用
                self.all_info[symbol]["adaptive_mean"] = mean_ret
                self.all_info[symbol]["adaptive_std"] = std_ret
                
                print(f"自适应阈值: mean={mean_ret:.6f}, std={std_ret:.6f}")
                print(f"标签分布: {np.bincount(y_train_list.astype(int) + 2)}")
                
            else:
                # 使用原有的离散化方法
                y_train_list = [self._discretize_return(r, return_scale, return_bins) for r in train_returns]
            
            X_train = np.array(X_train_list)
            y_train = np.array(y_train_list)
            
            # 检查是否有 NaN 或 Inf
            if np.any(np.isnan(X_train)) or np.any(np.isinf(X_train)):
                print(f"警告：训练数据包含 NaN 或 Inf，跳过训练")
                return None
            
            # 标准化特征
            X_train_scaled = self.all_info[symbol]["scaler"].fit_transform(X_train)
            
            # 训练 SVM 分类模型
            model = SVC(
                kernel='rbf',
                C=self.all_info[symbol].get("svm_C", 1.0),
                gamma=self.all_info[symbol].get("svm_gamma", 'scale'),
                class_weight='balanced'
            )
            model.fit(X_train_scaled, y_train)
            self.all_info[symbol]["model"] = model
            self.all_info[symbol]["last_train_index"] = current_index
            
            print(f"模型重训练完成 - Bar {current_index}, 训练样本数: {len(X_train)}")
        
        # 预测：使用最新的 lookback_period 期因子得分
        # 确保有足够的数据进行预测
        if (len(rsrs_scores) < lookback_period or 
            len(oc_scores) < lookback_period or 
            len(ret_scores) < lookback_period):
            return None
        
        # 确保模型已训练
        if self.all_info[symbol]["model"] is None:
            return None
        
        features_current = []
        features_current.extend(rsrs_scores[-lookback_period:])
        features_current.extend(oc_scores[-lookback_period:])
        features_current.extend(ret_scores[-lookback_period:])
        
        # 检查特征是否有 NaN 或 Inf
        if np.any(np.isnan(features_current)) or np.any(np.isinf(features_current)):
            print(f"警告：预测特征包含 NaN 或 Inf")
            return None
        
        X_current = np.array([features_current])
        X_current_scaled = self.all_info[symbol]["scaler"].transform(X_current)
        
        predicted_class = self.all_info[symbol]["model"].predict(X_current_scaled)[0]
        
        return predicted_class
    
    # ==================== 因子计算方法 ====================
    
    def _calculate_rsrs(self, symbol, kline):
        """计算 RSRS 因子得分"""
        rsrs_period = self.all_info[symbol].get("rsrs_period", 18)
        rsrs_std_period = self.all_info[symbol].get("rsrs_std_period", 600)
        
        high_prices = kline.high.values
        low_prices = kline.low.values
        rsrs_slopes = self.all_info[symbol]["rsrs_slopes"]
        
        current_len = len(high_prices)
        
        # 批量初始化
        if len(rsrs_slopes) < current_len - rsrs_period:
            start_idx = max(rsrs_period, len(rsrs_slopes))
            for i in range(start_idx, min(current_len, rsrs_std_period + rsrs_period)):
                window_high = high_prices[i-rsrs_period:i]
                window_low = low_prices[i-rsrs_period:i]
                X = window_low.reshape(-1, 1)
                y = window_high
                reg = LinearRegression().fit(X, y)
                rsrs_slopes.append(reg.coef_[0])
        
        # 增量更新
        if len(rsrs_slopes) >= rsrs_std_period:
            window_high = high_prices[-rsrs_period:]
            window_low = low_prices[-rsrs_period:]
            X = window_low.reshape(-1, 1)
            y = window_high
            reg = LinearRegression().fit(X, y)
            slope = reg.coef_[0]
            rsrs_slopes[:-1] = rsrs_slopes[1:]
            rsrs_slopes[-1] = slope
        
        if len(rsrs_slopes) == 0:
            return None
        
        # 计算标准化得分
        current_slope = rsrs_slopes[-1]
        recent_slopes = rsrs_slopes[-rsrs_std_period:]
        mean_slope = np.mean(recent_slopes)
        std_slope = np.std(recent_slopes)
        
        if std_slope == 0:
            return 0.0
        
        return (current_slope - mean_slope) / std_slope
    
    def _calculate_oc_corr(self, symbol, kline):
        """计算 OC 相关性因子得分"""
        score_period = self.all_info[symbol].get("oc_score_period", 20)
        std_period = self.all_info[symbol].get("oc_std_period", 600)
        
        close_prices = kline.close.values
        open_prices = kline.open.values
        corr_series = self.all_info[symbol]["corr_series"]
        
        current_len = len(close_prices)
        
        # 批量初始化
        if len(corr_series) < current_len - score_period:
            start_idx = max(score_period, len(corr_series))
            for i in range(start_idx, min(current_len, std_period + score_period)):
                window_close = close_prices[i-score_period:i]
                window_open = open_prices[i-score_period:i]
                mem = (window_close - window_open).mean()
                corr = np.corrcoef(window_close, window_open)[0][1] * (1 if mem >= 0 else -1)
                corr_series.append(corr)
        
        # 增量更新
        if len(corr_series) >= std_period:
            window_close = close_prices[-score_period:]
            window_open = open_prices[-score_period:]
            mem = (window_close - window_open).mean()
            corr = np.corrcoef(window_close, window_open)[0][1] * (1 if mem >= 0 else -1)
            corr_series[:-1] = corr_series[1:]
            corr_series[-1] = corr
        
        if len(corr_series) == 0:
            return None
        
        # 计算标准化得分
        current_score = corr_series[-1]
        recent_score = corr_series[-std_period:]
        mean_score = np.mean(recent_score)
        std_score = np.std(recent_score)
        
        if std_score == 0:
            return 0.0
        
        return (current_score - mean_score) / std_score
    
    def _calculate_ret(self, symbol, kline):
        """计算 RET 因子得分"""
        open_prices = kline.open.values
        high_prices = kline.high.values
        low_prices = kline.low.values
        close_prices = kline.close.values
        
        ret_series = self.all_info[symbol]["ret_series"]
        
        # 批量初始化
        if len(ret_series) == 0:
            ret_series.extend((high_prices + low_prices - open_prices - close_prices).tolist())
        else:
            # 增量更新
            ret_value = high_prices[-1] + low_prices[-1] - open_prices[-1] - close_prices[-1]
            ret_series.append(ret_value)
            if len(ret_series) > 800:
                ret_series.pop(0)
        
        if len(ret_series) == 0:
            return None
        
        # 返回当前值（不做标准化，因为 RET 本身就是一个相对值）
        return ret_series[-1]
